"""
模型训练配置
"""
import os
from typing import Dict, List, Optional, Union
import torch
import yaml
import datetime


class TrainingConfig:
    """
    训练配置类
    """
    def __init__(self, config_dict: Optional[Dict] = None):
        """
        初始化配置
        
        Args:
            config_dict: 配置字典
        """
        # 基础配置
        self.experiment_name = config_dict.get('experiment_name', 
                                             f"experiment_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}")
        self.project_root = config_dict.get('project_root', os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
        self.output_dir = config_dict.get('output_dir', os.path.join(self.project_root, 'outputs'))
        self.checkpoint_dir = config_dict.get('checkpoint_dir', os.path.join(self.output_dir, 'checkpoints'))
        self.log_dir = config_dict.get('log_dir', os.path.join(self.output_dir, 'logs'))
        self.data_dir = config_dict.get('data_dir', os.path.join(self.project_root, 'data'))
        
        # 设备配置
        self.use_gpu = config_dict.get('use_gpu', True)
        self.device = torch.device('cuda' if self.use_gpu and torch.cuda.is_available() else 'cpu')
        self.gpu_ids = config_dict.get('gpu_ids', [0])  # 使用的GPU ID列表
        
        # 数据配置
        self.train_data_path = config_dict.get('train_data_path', os.path.join(self.data_dir, 'train'))
        self.val_data_path = config_dict.get('val_data_path', os.path.join(self.data_dir, 'val'))
        self.test_data_path = config_dict.get('test_data_path', os.path.join(self.data_dir, 'test'))
        self.batch_size = config_dict.get('batch_size', 64)
        self.num_workers = config_dict.get('num_workers', 4)
        self.load_in_memory = config_dict.get('load_in_memory', True)
        self.use_balanced_sampler = config_dict.get('use_balanced_sampler', False)
        
        # 模型配置
        self.model_type = config_dict.get('model_type', 'policy_network')  # policy_network, value_network, actor_critic
        self.model_config = config_dict.get('model_config', {})
        
        # 训练配置
        self.learning_rate = config_dict.get('learning_rate', 1e-4)
        self.weight_decay = config_dict.get('weight_decay', 1e-5)
        self.optimizer_type = config_dict.get('optimizer_type', 'adam')  # adam, adamw, sgd
        self.scheduler_type = config_dict.get('scheduler_type', 'cosine')  # cosine, step, plateau
        self.scheduler_config = config_dict.get('scheduler_config', {
            'T_max': 100,  # cosine调度器的最大迭代次数
            'step_size': 30,  # step调度器的步长
            'gamma': 0.1,  # step调度器的衰减率
            'patience': 10  # plateau调度器的耐心值
        })
        
        self.epochs = config_dict.get('epochs', 100)
        self.max_steps = config_dict.get('max_steps', None)  # 最大步数，None表示使用epochs
        self.gradient_accumulation_steps = config_dict.get('gradient_accumulation_steps', 1)
        self.clip_grad_norm = config_dict.get('clip_grad_norm', 1.0)
        
        # 损失函数配置
        self.loss_type = config_dict.get('loss_type', 'cross_entropy')  # cross_entropy, mse, huber
        self.loss_config = config_dict.get('loss_config', {})
        
        # 早停配置
        self.early_stopping = config_dict.get('early_stopping', True)
        self.early_stopping_patience = config_dict.get('early_stopping_patience', 20)
        self.early_stopping_monitor = config_dict.get('early_stopping_monitor', 'val_loss')
        self.early_stopping_mode = config_dict.get('early_stopping_mode', 'min')  # min or max
        
        # 检查点配置
        self.save_checkpoint = config_dict.get('save_checkpoint', True)
        self.save_best_only = config_dict.get('save_best_only', True)
        self.checkpoint_interval = config_dict.get('checkpoint_interval', 1)  # 每多少个epoch保存一次
        self.max_checkpoints = config_dict.get('max_checkpoints', 5)  # 最多保存多少个检查点
        
        # 日志配置
        self.log_interval = config_dict.get('log_interval', 100)  # 每多少个batch记录一次日志
        self.use_tensorboard = config_dict.get('use_tensorboard', True)
        self.use_wandb = config_dict.get('use_wandb', False)
        self.wandb_project = config_dict.get('wandb_project', 'sichuan-mahjong-ai')
        
        # 评估配置
        self.evaluate_interval = config_dict.get('evaluate_interval', 1)  # 每多少个epoch评估一次
        self.eval_metrics = config_dict.get('eval_metrics', ['accuracy', 'f1'])
        
        # 其他配置
        self.seed = config_dict.get('seed', 42)
        self.resume_from_checkpoint = config_dict.get('resume_from_checkpoint', None)
        self.pretrained_model = config_dict.get('pretrained_model', None)
        self.freeze_layers = config_dict.get('freeze_layers', False)
        
        # 根据模型类型设置默认配置
        self._set_default_model_config()
        
        # 创建目录
        self._create_directories()
    
    def _set_default_model_config(self) -> None:
        """
        根据模型类型设置默认配置
        """
        if self.model_type == 'policy_network':
            self.model_config.setdefault('input_dim', 136)  # 4种花色 × 9个数字 × 4张牌 + 4种风牌 × 4 + 3种箭牌 × 4
            self.model_config.setdefault('hidden_dims', [512, 256, 128])
            self.model_config.setdefault('output_dim', 34)  # 34种可能的牌选择
            self.model_config.setdefault('dropout_rate', 0.3)
            self.model_config.setdefault('activation', 'relu')
            self.model_config.setdefault('use_batch_norm', True)
        
        elif self.model_type == 'value_network':
            self.model_config.setdefault('input_dim', 136)
            self.model_config.setdefault('hidden_dims', [512, 256, 128])
            self.model_config.setdefault('output_dim', 1)  # 单个值输出
            self.model_config.setdefault('dropout_rate', 0.3)
            self.model_config.setdefault('activation', 'relu')
            self.model_config.setdefault('use_batch_norm', True)
        
        elif self.model_type == 'actor_critic':
            self.model_config.setdefault('input_dim', 136)
            self.model_config.setdefault('shared_hidden_dims', [512, 256])
            self.model_config.setdefault('actor_hidden_dims', [128])
            self.model_config.setdefault('critic_hidden_dims', [128])
            self.model_config.setdefault('actor_output_dim', 34)
            self.model_config.setdefault('critic_output_dim', 1)
            self.model_config.setdefault('dropout_rate', 0.3)
            self.model_config.setdefault('activation', 'relu')
        
        elif self.model_type == 'dqn':
            self.model_config.setdefault('input_dim', 136)
            self.model_config.setdefault('hidden_dims', [512, 256, 128])
            self.model_config.setdefault('output_dim', 34)
            self.model_config.setdefault('gamma', 0.99)  # 折扣因子
            self.model_config.setdefault('epsilon_start', 1.0)  # 探索率初始值
            self.model_config.setdefault('epsilon_end', 0.1)  # 探索率最小值
            self.model_config.setdefault('epsilon_decay', 0.995)  # 探索率衰减率
            self.model_config.setdefault('target_update_freq', 10)  # 目标网络更新频率
    
    def _create_directories(self) -> None:
        """
        创建必要的目录
        """
        os.makedirs(self.output_dir, exist_ok=True)
        os.makedirs(self.checkpoint_dir, exist_ok=True)
        os.makedirs(self.log_dir, exist_ok=True)
        os.makedirs(os.path.join(self.log_dir, self.experiment_name), exist_ok=True)
    
    def to_dict(self) -> Dict:
        """
        转换为字典
        
        Returns:
            配置字典
        """
        return {
            'experiment_name': self.experiment_name,
            'project_root': self.project_root,
            'output_dir': self.output_dir,
            'checkpoint_dir': self.checkpoint_dir,
            'log_dir': self.log_dir,
            'data_dir': self.data_dir,
            'use_gpu': self.use_gpu,
            'device': str(self.device),
            'gpu_ids': self.gpu_ids,
            'train_data_path': self.train_data_path,
            'val_data_path': self.val_data_path,
            'test_data_path': self.test_data_path,
            'batch_size': self.batch_size,
            'num_workers': self.num_workers,
            'load_in_memory': self.load_in_memory,
            'use_balanced_sampler': self.use_balanced_sampler,
            'model_type': self.model_type,
            'model_config': self.model_config,
            'learning_rate': self.learning_rate,
            'weight_decay': self.weight_decay,
            'optimizer_type': self.optimizer_type,
            'scheduler_type': self.scheduler_type,
            'scheduler_config': self.scheduler_config,
            'epochs': self.epochs,
            'max_steps': self.max_steps,
            'gradient_accumulation_steps': self.gradient_accumulation_steps,
            'clip_grad_norm': self.clip_grad_norm,
            'loss_type': self.loss_type,
            'loss_config': self.loss_config,
            'early_stopping': self.early_stopping,
            'early_stopping_patience': self.early_stopping_patience,
            'early_stopping_monitor': self.early_stopping_monitor,
            'early_stopping_mode': self.early_stopping_mode,
            'save_checkpoint': self.save_checkpoint,
            'save_best_only': self.save_best_only,
            'checkpoint_interval': self.checkpoint_interval,
            'max_checkpoints': self.max_checkpoints,
            'log_interval': self.log_interval,
            'use_tensorboard': self.use_tensorboard,
            'use_wandb': self.use_wandb,
            'wandb_project': self.wandb_project,
            'evaluate_interval': self.evaluate_interval,
            'eval_metrics': self.eval_metrics,
            'seed': self.seed,
            'resume_from_checkpoint': self.resume_from_checkpoint,
            'pretrained_model': self.pretrained_model,
            'freeze_layers': self.freeze_layers
        }
    
    def save(self, file_path: str) -> None:
        """
        保存配置到文件
        
        Args:
            file_path: 文件路径
        """
        with open(file_path, 'w', encoding='utf-8') as f:
            yaml.dump(self.to_dict(), f, default_flow_style=False, allow_unicode=True)
    
    @classmethod
    def load(cls, file_path: str) -> 'TrainingConfig':
        """
        从文件加载配置
        
        Args:
            file_path: 文件路径
            
        Returns:
            配置实例
        """
        with open(file_path, 'r', encoding='utf-8') as f:
            config_dict = yaml.safe_load(f)
        
        return cls(config_dict)
    
    def update(self, **kwargs) -> 'TrainingConfig':
        """
        更新配置
        
        Args:
            **kwargs: 要更新的配置项
            
        Returns:
            更新后的配置实例
        """
        config_dict = self.to_dict()
        config_dict.update(kwargs)
        
        return TrainingConfig(config_dict)


# 预设配置
PRESET_CONFIGS = {
    'default': {
        'model_type': 'policy_network',
        'batch_size': 64,
        'learning_rate': 1e-4,
        'epochs': 100,
        'use_gpu': True
    },
    
    'policy_network': {
        'model_type': 'policy_network',
        'model_config': {
            'input_dim': 136,
            'hidden_dims': [512, 256, 128],
            'output_dim': 34,
            'dropout_rate': 0.3
        }
    },
    
    'value_network': {
        'model_type': 'value_network',
        'model_config': {
            'input_dim': 136,
            'hidden_dims': [512, 256, 128],
            'output_dim': 1,
            'dropout_rate': 0.3
        }
    },
    
    'actor_critic': {
        'model_type': 'actor_critic',
        'model_config': {
            'input_dim': 136,
            'shared_hidden_dims': [512, 256],
            'actor_hidden_dims': [128],
            'critic_hidden_dims': [128],
            'dropout_rate': 0.3
        }
    },
    
    'dqn': {
        'model_type': 'dqn',
        'model_config': {
            'input_dim': 136,
            'hidden_dims': [512, 256, 128],
            'output_dim': 34,
            'gamma': 0.99,
            'epsilon_start': 1.0,
            'epsilon_end': 0.1,
            'epsilon_decay': 0.995
        }
    },
    
    'fast_dev': {
        'epochs': 3,
        'batch_size': 16,
        'num_workers': 0,
        'log_interval': 10,
        'checkpoint_interval': 1
    }
}


def get_preset_config(preset_name: str, **kwargs) -> TrainingConfig:
    """
    获取预设配置
    
    Args:
        preset_name: 预设名称
        **kwargs: 额外的配置项
        
    Returns:
        配置实例
    """
    # 获取预设配置
    if preset_name not in PRESET_CONFIGS:
        raise ValueError(f"未知的预设配置: {preset_name}")
    
    # 合并配置
    config_dict = {'experiment_name': preset_name}
    
    # 先添加默认配置
    if 'default' in PRESET_CONFIGS:
        config_dict.update(PRESET_CONFIGS['default'])
    
    # 再添加指定预设
    config_dict.update(PRESET_CONFIGS[preset_name])
    
    # 最后添加自定义配置
    config_dict.update(kwargs)
    
    return TrainingConfig(config_dict)


def create_config_from_args(args) -> TrainingConfig:
    """
    从命令行参数创建配置
    
    Args:
        args: 命令行参数
        
    Returns:
        配置实例
    """
    config_dict = {}
    
    # 基础配置
    if hasattr(args, 'experiment_name'):
        config_dict['experiment_name'] = args.experiment_name
    if hasattr(args, 'output_dir'):
        config_dict['output_dir'] = args.output_dir
    if hasattr(args, 'data_dir'):
        config_dict['data_dir'] = args.data_dir
    
    # 训练配置
    if hasattr(args, 'epochs'):
        config_dict['epochs'] = args.epochs
    if hasattr(args, 'batch_size'):
        config_dict['batch_size'] = args.batch_size
    if hasattr(args, 'learning_rate'):
        config_dict['learning_rate'] = args.learning_rate
    
    # 模型配置
    if hasattr(args, 'model_type'):
        config_dict['model_type'] = args.model_type
    
    # 其他配置
    if hasattr(args, 'use_gpu'):
        config_dict['use_gpu'] = args.use_gpu
    if hasattr(args, 'seed'):
        config_dict['seed'] = args.seed
    if hasattr(args, 'resume_from_checkpoint'):
        config_dict['resume_from_checkpoint'] = args.resume_from_checkpoint
    if hasattr(args, 'pretrained_model'):
        config_dict['pretrained_model'] = args.pretrained_model
    
    return TrainingConfig(config_dict)


def print_config_summary(config: TrainingConfig) -> None:
    """
    打印配置摘要
    
    Args:
        config: 配置实例
    """
    print("=" * 60)
    print("训练配置摘要")
    print("=" * 60)
    
    # 实验信息
    print(f"实验名称: {config.experiment_name}")
    print(f"输出目录: {config.output_dir}")
    print(f"检查点目录: {config.checkpoint_dir}")
    print(f"日志目录: {config.log_dir}")
    
    # 设备信息
    print(f"设备: {config.device}")
    print(f"使用GPU: {config.use_gpu}")
    
    # 数据信息
    print(f"训练数据路径: {config.train_data_path}")
    print(f"验证数据路径: {config.val_data_path}")
    print(f"测试数据路径: {config.test_data_path}")
    print(f"批量大小: {config.batch_size}")
    print(f"工作线程数: {config.num_workers}")
    
    # 模型信息
    print(f"模型类型: {config.model_type}")
    print(f"模型配置: {config.model_config}")
    
    # 训练信息
    print(f"学习率: {config.learning_rate}")
    print(f"权重衰减: {config.weight_decay}")
    print(f"优化器: {config.optimizer_type}")
    print(f"学习率调度器: {config.scheduler_type}")
    print(f"训练轮数: {config.epochs}")
    print(f"梯度累积步数: {config.gradient_accumulation_steps}")
    
    # 早停和检查点
    print(f"早停: {'开启' if config.early_stopping else '关闭'}")
    if config.early_stopping:
        print(f"早停耐心值: {config.early_stopping_patience}")
    print(f"保存检查点: {'开启' if config.save_checkpoint else '关闭'}")
    print(f"仅保存最佳模型: {'是' if config.save_best_only else '否'}")
    
    print("=" * 60)